import os

import cv2 as cv
import numpy as np
import struct


dataset_dir = 'data'
train_data = []
train_label = []

for label_name in os.listdir(dataset_dir):
    data_dir = os.path.join(dataset_dir, label_name)
    for image_name in os.listdir(data_dir):
        image_path = os.path.join(data_dir, image_name)
        image = cv.imread(image_path, 0)
        image = cv.resize(image,(80, 80))
        train_data.append(image)
        train_label.append(label_name)

train_data = np.array(train_data, dtype=np.float32)
train_label = np.array(train_label, dtype=np.float32)
train_data = train_data.reshape(train_data.shape[0], 6400)/255
knn = cv.ml.KNearest_create()

knn.setDefaultK(5)

knn.setIsClassifier(True)

knn.train(train_data, cv.ml.ROW_SAMPLE, train_label)

knn.save('block_knn2.xml')


# train_label = np.array(train_label, dtype=np.float32)





